import pandas as pd 
import numpy as np 
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, classification_report
from joblib import dump

class ClassificationModel(object):
    def __init__(self):
        # 读取数据时增加异常处理，避免文件不存在或格式错误
        try:
            self.df = pd.read_csv('fina_indicator.csv', encoding='utf-8-sig')
            print(f"成功读取数据，共 {len(self.df)} 条记录")
        except FileNotFoundError:
            raise FileNotFoundError("错误：未找到 fina_indicator.csv 文件，请先运行数据提取脚本生成该文件")
        except Exception as e:
            raise Exception(f"读取数据失败：{str(e)}")
    
    def get_conditions(self):
        df = self.df
        
        # 数据校验：确保必要的列存在
        required_cols = ['max_closes', 'min_closes', 'the_closes', 'eps', 'total_revenue_ps', 
                        'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe', 'tangible_asset',
                        'bps', 'grossprofit_margin', 'npta']
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            raise KeyError(f"数据缺少必要列：{', '.join(missing_cols)}")
        
        # 计算收益潜力和风险程度（避免除零错误）
        df['max_ratio'] = df['max_closes'] / df['the_closes'].replace(0, np.nan)
        df['min_ratio'] = df['min_closes'] / df['the_closes'].replace(0, np.nan)
        
        # 剔除计算失败的行（比值为 NaN 的情况）
        df = df.dropna(subset=['max_ratio', 'min_ratio']).reset_index(drop=True)
        if len(df) == 0:
            raise ValueError("没有有效数据用于计算收益/风险比值（可能 the_closes 全为 0）")
        
        # 划分高低阈值（前40%为高收益/高风险）
        high_return_threshold = df['max_ratio'].quantile(0.4)
        high_risk_threshold = df['min_ratio'].quantile(0.4)
        print(f"高收益阈值：{high_return_threshold:.4f}，高风险阈值：{high_risk_threshold:.4f}")
        
        # 生成分类标签
        conditions = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] >= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] < high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] < high_risk_threshold)
        ]
        labels = ['高收益高风险', '低收益高风险', '高收益低风险', '低收益低风险']
        df['category'] = np.select(conditions, labels, default='未知')
        
        # 打印分类分布
        print("\n分类标签分布：")
        print(df['category'].value_counts())
        
        # 标签编码
        label_encoder = LabelEncoder()
        df['category_encoded'] = label_encoder.fit_transform(df['category'])
        
        # 数据标准化
        scaler = StandardScaler()
        features = df[['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe',
                      'tangible_asset', 'bps', 'grossprofit_margin', 'npta']]
        X = scaler.fit_transform(features)
        y = df['category_encoded']
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=24, stratify=y)
        print(f"\n训练集大小：{X_train.shape}，测试集大小：{X_test.shape}")
        
        # 保存标准化器
        dump(scaler, 'scaler.joblib')
        print("标准化器已保存为 scaler.joblib")
        
        return X_train, X_test, y_train, y_test, label_encoder

    # 修正：将 knn_utils 缩进为类方法
    def knn_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        print("\n" + "="*50)
        print("开始训练 KNN 模型...")
        # 初始化 KNN 分类器（n_neighbors=5 为默认值，可根据需求调整）
        knn = KNeighborsClassifier(n_neighbors=5)
        # 训练模型
        knn.fit(X_train, y_train)
        # 预测测试集
        y_pred = knn.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"\nKNN 模型准确率: {accuracy:.4f}")
        print("KNN 分类报告:")
        print(report)
        # 保存模型
        dump(knn, 'knn.joblib')
        print("KNN 模型已保存为 knn.joblib")

    # 新增：实现 svc_utils 方法（SVM 分类）
    def svc_utils(self, X_train, X_test, y_train, y_test, label_encoder):
        print("\n" + "="*50)
        print("开始训练 SVM 模型...")
        # 初始化 SVM 分类器（默认 RBF 核，可根据需求调整参数）
        svc = SVC(kernel='rbf', random_state=24, probability=True)
        # 训练模型
        svc.fit(X_train, y_train)
        # 预测测试集
        y_pred = svc.predict(X_test)
        # 评估模型
        accuracy = accuracy_score(y_test, y_pred)
        report = classification_report(y_test, y_pred, target_names=label_encoder.classes_)
        print(f"\nSVM 模型准确率: {accuracy:.4f}")
        print("SVM 分类报告:")
        print(report)
        # 保存模型
        dump(svc, 'svc.joblib')
        print("SVM 模型已保存为 svc.joblib")

if __name__ == '__main__':
    try:
        # 初始化模型类
        cu = ClassificationModel()
        # 数据预处理和特征工程
        X_train, X_test, y_train, y_test, label_encoder = cu.get_conditions()
        # 训练 KNN 模型（取消注释即可运行）
        # cu.knn_utils(X_train, X_test, y_train, y_test, label_encoder)
        # 训练 SVM 模型（当前默认运行）
        cu.svc_utils(X_train, X_test, y_train, y_test, label_encoder)
        print("\n" + "="*50)
        print("程序运行完成！")
    except Exception as e:
        print(f"\n程序运行失败：{str(e)}")
